#include "common.hpp"

std::ofstream result;

int Task(Direction direction, size_t device_id, size_t data_size, int n_warmup, size_t n_loop) {
    ACLCHECK(aclrtSetDevice(device_id));

    aclrtStream stream;
    aclrtEvent totalEndEvent;
    aclrtEvent totalStartEvent;
    std::vector<aclrtEvent> endEvents(n_loop);
    std::vector<aclrtEvent> startEvents(n_loop);
    ACLCHECK(aclrtCreateStream(&stream));
    ACLCHECK(aclrtCreateEvent(&totalEndEvent));
    ACLCHECK(aclrtCreateEvent(&totalStartEvent));
    for (int i = 0; i < n_loop; i++) {
        ACLCHECK(aclrtCreateEvent(&endEvents[i]));
        ACLCHECK(aclrtCreateEvent(&startEvents[i]));
    }

    // 申请Host内存, 并初始化, 让每一块的数据都不一样
    float *hostBuff;
    float *deviceBuff;
    ACLCHECK(aclrtMallocHost((void **)&hostBuff, data_size));
    ACLCHECK(aclrtMalloc((void **)&deviceBuff, data_size, ACL_MEM_MALLOC_HUGE_FIRST));

    // 定义测试的Lambda函数
    auto inner_func = [data_size, stream, hostBuff, deviceBuff, startEvents, endEvents, direction](bool warm_up, size_t i) -> int {
        if (!warm_up) ACLCHECK(aclrtRecordEvent(startEvents[i], stream));
        if (direction == Direction::H2D) {
            ACLCHECK(aclrtMemcpyAsync(deviceBuff, data_size, hostBuff, data_size, ACL_MEMCPY_HOST_TO_DEVICE, stream));
        } else if (direction == Direction::D2H) {
            ACLCHECK(aclrtMemcpyAsync(hostBuff, data_size, deviceBuff, data_size, ACL_MEMCPY_DEVICE_TO_HOST, stream));
        } else {
            printf("error\n");
        }
        if (!warm_up) ACLCHECK(aclrtRecordEvent(endEvents[i], stream));
        return 0;
    };
    // Warmup阶段
    for (size_t i = 0; i < n_warmup; i++) inner_func(true, i);

    // 执行测试的代码
    ACLCHECK(aclrtRecordEvent(totalStartEvent, stream));
    for (size_t i = 0; i < n_loop; i++) inner_func(false, i);
    ACLCHECK(aclrtRecordEvent(totalEndEvent, stream));
    ACLCHECK(aclrtSynchronizeStream(stream));

    // 计算结果
    float total_time;
    ACLCHECK(aclrtEventElapsedTime(&total_time, totalStartEvent, totalEndEvent));

    // 打印标题信息
    result << data_size << ","         //
           << total_time * 1E3 << ","  //
           << (n_loop * (float)data_size / GB) / (total_time * 1E-3);
    std::vector<float> times(n_loop);
    for (size_t i = 0; i < n_loop; i++) {
        ACLCHECK(aclrtEventElapsedTime(&times[i], startEvents[i], endEvents[i]));
        result << "," << times[i] * 1E3;
    }
    for (size_t i = 0; i < n_loop; i++) result << "," << ((float)data_size / GB) / (times[i] * 1E-3);
    result << std::endl;
    result.flush();
    std::cout << (direction == Direction::H2D ? "H2D" : "D2H") << "_Task >> device: " << device_id               //
              << ", n_loop: " << n_loop << ", data_size: " << GetSizeStr(data_size) << ", time: " << total_time  //
              << " ms, bandwidth: " << GetBandwidthStr((n_loop * (float)data_size / GB) / (total_time * 1E-3)) << std::endl;

    // 清理数据
    ACLCHECK(aclrtFree(deviceBuff));
    ACLCHECK(aclrtFreeHost(hostBuff));

    ACLCHECK(aclrtDestroyEvent(totalEndEvent));
    ACLCHECK(aclrtDestroyEvent(totalStartEvent));
    for (size_t i = 0; i < n_loop; i++) {
        ACLCHECK(aclrtDestroyEvent(endEvents[i]));
        ACLCHECK(aclrtDestroyEvent(startEvents[i]));
    }
    ACLCHECK(aclrtDestroyStream(stream));
    ACLCHECK(aclrtResetDevice(device_id));

    return 0;
}

int main(int argc, char *argv[]) {
    size_t device_id = 7;                    // 获取设备的ID
    size_t n_warmup = 5;                     //
    size_t n_loop = 10;                     //
    size_t start_data_size = 8;              // 测试的起始数据包大小
    size_t end_data_size = (size_t)8 * GB;  // 测试的最终数据大小
    Direction direction = Direction::D2H;    //

    // 创建结果保存文件
    std::ostringstream filename;
    filename << "perf-single-stream-device" << device_id << "-loop" << n_loop << "-" << (direction == Direction::H2D ? "h2d" : "d2h") << ".csv";
    result.open(filename.str());

    // 写入标题
    result << "data_size(byte),total_time(us),total_bandwidth(GB/s)";
    for (size_t i = 0; i < n_loop; i++) result << ",time_" << i + 1 << "(us)";
    for (size_t i = 0; i < n_loop; i++) result << ",bandwidth_" << i + 1 << "(GB/s)";
    result << std::endl;

    // 开始测试
    size_t data_size = start_data_size;
    while (data_size <= end_data_size) {
        Task(direction, device_id, data_size, n_warmup, n_loop);
        data_size *= 2;
    }
    result.flush();
    result.close();
    return 0;

    return 0;
}
